package cn.ibizlab.core.oauth2.service.impl;

import cn.ibizlab.core.authentication.domain.AuthGrantedRole;
import cn.ibizlab.core.authentication.domain.AuthGrantedSystem;
import cn.ibizlab.core.authentication.domain.AuthInfo;
import cn.ibizlab.core.oauth2.domain.OAuth2TokenPayload;
import cn.ibizlab.core.oauth2.mapping.OAuth2TokenPayloadMapping;
import cn.ibizlab.util.enums.Entities;
import cn.ibizlab.util.errors.BadRequestAlertException;
import cn.ibizlab.util.security.AuthUser20Impl;
import cn.ibizlab.util.security.JwkStore;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jose.JWSSigner;
import com.nimbusds.jose.crypto.MACSigner;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jwt.JWTClaimNames;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import lombok.Data;
import lombok.SneakyThrows;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.oauth2.common.*;
import org.springframework.security.oauth2.common.exceptions.InvalidGrantException;
import org.springframework.security.oauth2.common.exceptions.InvalidScopeException;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.OAuth2Request;
import org.springframework.security.oauth2.provider.TokenRequest;
import org.springframework.security.oauth2.provider.token.*;
import org.springframework.security.web.authentication.preauth.PreAuthenticatedAuthenticationToken;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.ObjectUtils;

import java.security.PrivateKey;
import java.util.*;

/**
 * OAuth2令牌服务：将原有OAuth2Token(uuid)调整为由 nimbus-jose-jwt 生成访问令牌，访问令牌payload中存储用户基本信息
 */
@Data
public class OAuth2TokenService extends DefaultTokenServices {

    protected static final String SIGNATURE_RSA = "RSA";
    protected static final String SIGNATURE_MAC = "MAC";
    private ObjectMapper mapper = new ObjectMapper();
    private TokenStore tokenStore;
    private TokenEnhancer accessTokenEnhancer;

    @Value("${ibiz.jwt.signature:MAC}")
    private String signature;
    @Value("${ibiz.jwt.secret:ibzsecret}")
    private String secret;
    @Autowired
    OAuth2TokenPayloadMapping oAuth2TokenPayloadMapping;
    private boolean supportRefreshToken = false;
    private boolean reuseRefreshToken = true;
    private AuthenticationManager authenticationManager;


    @Autowired
    @Lazy
    private JwkStore jwkStore;

    private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();

    private static final String GLOBAL_ROLE = "_global";

    private static final String ROLE_API_USERS = "ROLE_APIUSERS";

    private static final String ROLE_SUPER_ADMIN = "ROLE_SUPERADMIN";

    @Override
    public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
        try {
            return tokenStore.getAccessToken(authentication);
        }catch (Exception ex) {
            return null;
        }
    }

    @Value("${ibiz.auth.tokenExclusion:false}")
    private boolean tokenExclusion;


    @Transactional
    @Override
    public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication) throws AuthenticationException {
        OAuth2RefreshToken refreshToken = null;

        if(tokenExclusion) {
            OAuth2AccessToken existingAccessToken = tokenStore.getAccessToken(authentication);
            if (existingAccessToken != null) {
                if (existingAccessToken.isExpired()) {
                    if (existingAccessToken.getRefreshToken() != null) {
                        refreshToken = existingAccessToken.getRefreshToken();
                        // The token store could remove the refresh token when the
                        // access token is removed, but we want to
                        // be sure...
                        tokenStore.removeRefreshToken(refreshToken);
                    }
                    tokenStore.removeAccessToken(existingAccessToken);
                } else {
                    // Re-store the access token in case the authentication has changed
                    tokenStore.storeAccessToken(existingAccessToken, authentication);
                    return existingAccessToken;
                }
            }
        }

        // Only create a new refresh token if there wasn't an existing one
        // associated with an expired access token.
        // Clients might be holding existing refresh tokens, so we re-use it in
        // the case that the old access token
        // expired.
        if (refreshToken == null) {
            refreshToken = createRefreshToken(authentication);
        }
        // But the refresh token itself might need to be re-issued if it has
        // expired.
        else if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
            ExpiringOAuth2RefreshToken expiring = (ExpiringOAuth2RefreshToken) refreshToken;
            if (System.currentTimeMillis() > expiring.getExpiration().getTime()) {
                refreshToken = createRefreshToken(authentication);
            }
        }

        OAuth2AccessToken accessToken = createAccessToken(authentication, refreshToken);
        accessToken.getAdditionalInformation().put("token_id",
                authenticationKeyGenerator.extractKey(authentication));
        if("client_credentials".equals(authentication.getOAuth2Request().getGrantType())) {
            AuthGrantedSystem system = new AuthGrantedSystem().setDeploySystemId(GLOBAL_ROLE).setSystemId(GLOBAL_ROLE).setDisplayName("全局").setRoles(new ArrayList<>());
            if(!ObjectUtils.isEmpty(authentication.getAuthorities())){
                for(GrantedAuthority authority : authentication.getAuthorities()) {
                    system.getRoles().add(new AuthGrantedRole().setRoleId(authority.getAuthority()).setAuthority(authority.getAuthority()).setDisplayName(authority.getAuthority()));
                }
            }
            AuthInfo authInfo = new AuthInfo().setId(authentication.getOAuth2Request().getClientId()).setUsername(authentication.getOAuth2Request().getClientId()).setDisplayName(authentication.getOAuth2Request().getClientId()).setGrants(Arrays.asList(system));
            accessToken.getAdditionalInformation().put("user",authInfo);
        }

        tokenStore.storeAccessToken(accessToken, authentication);
        // In case it was modified
        refreshToken = accessToken.getRefreshToken();
        if (refreshToken != null) {
            tokenStore.storeRefreshToken(refreshToken, authentication);
        }

        return accessToken;

    }

    @Override
    @Transactional(noRollbackFor={InvalidTokenException.class, InvalidGrantException.class})
    public OAuth2AccessToken refreshAccessToken(String refreshTokenValue, TokenRequest tokenRequest)
            throws AuthenticationException {

        if (!supportRefreshToken) {
            throw new InvalidGrantException("Invalid refresh token: " + refreshTokenValue);
        }

        OAuth2RefreshToken refreshToken = tokenStore.readRefreshToken(refreshTokenValue);
        if (refreshToken == null) {
            throw new InvalidGrantException("Invalid refresh token: " + refreshTokenValue);
        }

        OAuth2Authentication authentication = tokenStore.readAuthenticationForRefreshToken(refreshToken);
        if (this.authenticationManager != null && !authentication.isClientOnly()) {
            // The client has already been authenticated, but the user authentication might be old now, so give it a
            // chance to re-authenticate.
            Authentication user = new PreAuthenticatedAuthenticationToken(authentication.getUserAuthentication(), "", authentication.getAuthorities());
            user = authenticationManager.authenticate(user);
            Object details = authentication.getDetails();
            authentication = new OAuth2Authentication(authentication.getOAuth2Request(), user);
            authentication.setDetails(details);
        }
        String clientId = authentication.getOAuth2Request().getClientId();
        if (clientId == null || !clientId.equals(tokenRequest.getClientId())) {
            throw new InvalidGrantException("Wrong client for this refresh token: " + refreshTokenValue);
        }

        // clear out any access tokens already associated with the refresh
        // token.
        tokenStore.removeAccessTokenUsingRefreshToken(refreshToken);

        if (isExpired(refreshToken)) {
            tokenStore.removeRefreshToken(refreshToken);
            throw new InvalidTokenException("Invalid refresh token (expired): " + refreshToken);
        }

        authentication = createRefreshedAuthentication(authentication, tokenRequest);

        if (!reuseRefreshToken) {
            tokenStore.removeRefreshToken(refreshToken);
            refreshToken = createRefreshToken(authentication);
        }

        OAuth2AccessToken accessToken = createAccessToken(authentication, refreshToken);
        tokenStore.storeAccessToken(accessToken, authentication);
        if (!reuseRefreshToken) {
            tokenStore.storeRefreshToken(accessToken.getRefreshToken(), authentication);
        }
        return accessToken;
    }

    /**
     * Create a refreshed authentication.
     *
     * @param authentication The authentication.
     * @param request The scope for the refreshed token.
     * @return The refreshed authentication.
     * @throws InvalidScopeException If the scope requested is invalid or wider than the original scope.
     */
    private OAuth2Authentication createRefreshedAuthentication(OAuth2Authentication authentication, TokenRequest request) {
        Set<String> scope = request.getScope();
        OAuth2Request clientAuth = authentication.getOAuth2Request().refresh(request);
        if (scope != null && !scope.isEmpty()) {
            Set<String> originalScope = clientAuth.getScope();
            if (originalScope == null || !originalScope.containsAll(scope)) {
                throw new InvalidScopeException("Unable to narrow the scope of the client authentication to " + scope
                        + ".", originalScope);
            }
            else {
                clientAuth = clientAuth.narrowScope(scope);
            }
        }
        return new OAuth2Authentication(clientAuth, authentication.getUserAuthentication());
    }

	private OAuth2RefreshToken createRefreshToken(OAuth2Authentication authentication) {
		if (!isSupportRefreshToken(authentication.getOAuth2Request())) {
			return null;
		}
		int validitySeconds = getRefreshTokenValiditySeconds(authentication.getOAuth2Request());
		String value = UUID.randomUUID().toString();
		if (validitySeconds > 0) {
			return new DefaultExpiringOAuth2RefreshToken(value, new Date(System.currentTimeMillis()
					+ (validitySeconds * 1000L)));
		}
		return new DefaultOAuth2RefreshToken(value);
	}

    /**
     *  nimbus-jose-jwt 生成访问令牌，访问令牌payload中存储用户基本信息
     * @param authentication
     * @param refreshToken
     * @return
     */
    private OAuth2AccessToken createAccessToken(OAuth2Authentication authentication, OAuth2RefreshToken refreshToken) {
        DefaultOAuth2AccessToken token = new DefaultOAuth2AccessToken(UUID.randomUUID().toString());
        int validitySeconds = getAccessTokenValiditySeconds(authentication.getOAuth2Request());
        if (validitySeconds > 0) {
            token.setExpiration(new Date(System.currentTimeMillis() + (validitySeconds * 1000L)));
        }
        token.setAdditionalInformation(new HashMap<>());
        //通过 nimbus-jose-jwt 生成访问令牌
        token.setValue(getAccessToken(authentication,token.getExpiration()));
        token.setRefreshToken(refreshToken);
        token.setScope(authentication.getOAuth2Request().getScope());
        return accessTokenEnhancer != null ? accessTokenEnhancer.enhance(token, authentication) : token;
    }


    /**
     * 生成 accessToken
     * @param authentication
     * @param expiration
     * @return
     */
    protected String getAccessToken(OAuth2Authentication authentication,Date expiration){
        Map<String,Object> tokenPayload = getAccessTokenPayload(authentication,expiration);
        if(ObjectUtils.isEmpty(signature)|| SIGNATURE_MAC.equalsIgnoreCase(signature))
            return generateTokenByMAC(tokenPayload);
        else if(SIGNATURE_RSA.equalsIgnoreCase(signature))
            return generateTokenByRSA(tokenPayload);
        else
            throw new BadRequestAlertException(String.format("生成访问令牌出错，签名暂未支持[%1$s]加密算法",signature), Entities.AUTH_USER.toString(),authentication.getName());
    }

    /**
     * 通过MAC方式生成token（对称加密）
     * @param tokenPayload
     * @return
     */
    @SneakyThrows
    protected String generateTokenByMAC(Map<String,Object> tokenPayload){
        JWSSigner signer = new MACSigner(secret);
        JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.HS512)
                .contentType("JWS")
                .build();

        JWTClaimsSet claimsSet = JWTClaimsSet.parse(tokenPayload);
        SignedJWT signedJWT = new SignedJWT(header,claimsSet);
        signedJWT.sign(signer);
        return signedJWT.serialize();
    }

    /**
     * 通过RSA方式生成token（非对称加密）
     * @param tokenPayload
     * @return
     */
    @SneakyThrows
    protected String generateTokenByRSA(Map<String,Object> tokenPayload) {
        PrivateKey privateJWK = jwkStore.getPrivateKey();
        JWSSigner signer = new RSASSASigner(privateJWK);
        JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.RS256)
                .contentType("JWS")
                .build();

        JWTClaimsSet claimsSet = JWTClaimsSet.parse(tokenPayload);
        SignedJWT signedJWT = new SignedJWT(header,claimsSet);
        signedJWT.sign(signer);
        return signedJWT.serialize();
    }


    /**
     * 构造token载荷体
     * @param authentication
     * @return
     */
    protected Map<String, Object> getAccessTokenPayload(Authentication authentication,Date expiration){
        AuthUser20Impl authenticationUser = null;
        Map<String, Object> payload = new HashMap<>();
        if(authentication.getDetails() instanceof AuthUser20Impl){//常规登录模式，获取用户信息存入token
           authenticationUser = (AuthUser20Impl) authentication.getDetails();
        }
        else if(authentication instanceof OAuth2Authentication){ //oauth2 password模式，获取用户信息存入token
            OAuth2Authentication oAuth2Authentication = (OAuth2Authentication) authentication;
            if(oAuth2Authentication.getUserAuthentication() != null && oAuth2Authentication.getUserAuthentication().getPrincipal() !=null
                                                                    && oAuth2Authentication.getUserAuthentication().getPrincipal() instanceof AuthUser20Impl){
                authenticationUser = (AuthUser20Impl) oAuth2Authentication.getUserAuthentication().getPrincipal();
            }
        }
        if(authenticationUser != null){
            OAuth2TokenPayload tokenPayload = oAuth2TokenPayloadMapping.toDto(authenticationUser);
            payload = mapper.convertValue(tokenPayload, new TypeReference<Map<String, Object>>() {});
        }
        payload.put(JWTClaimNames.SUBJECT,authentication.getName());
        payload.put(JWTClaimNames.ISSUED_AT,(System.currentTimeMillis())/1000);
        payload.put(JWTClaimNames.JWT_ID,UUID.randomUUID().toString());
        payload.put(JWTClaimNames.EXPIRATION_TIME, (!ObjectUtils.isEmpty(expiration) ? (expiration.getTime())/1000: expiration));
        return payload;

    }

    @Override
    public void setTokenStore(TokenStore tokenStore) {
        this.tokenStore = tokenStore;
        super.setTokenStore(tokenStore);
    }

    @Override
    public void setAuthenticationManager(AuthenticationManager authenticationManager) {
        this.authenticationManager = authenticationManager;
        super.setAuthenticationManager(authenticationManager);
    }

    @Override
    public void setSupportRefreshToken(boolean supportRefreshToken) {
        this.supportRefreshToken = supportRefreshToken;
        super.setSupportRefreshToken(supportRefreshToken);
    }

    @Override
    public void setReuseRefreshToken(boolean reuseRefreshToken) {
        this.reuseRefreshToken = reuseRefreshToken;
        super.setReuseRefreshToken(reuseRefreshToken);
    }
}
